"""
<<Deep Learning With Pytorch>>,Chapter 2.2 引例
                            ---变马为斑马
Copied by:MIMIC     2024/6/15 12:27
"""
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt


def build_conv_block(dim):
    conv_block = []

    conv_block += [nn.ReflectionPad2d(1)]
    conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=0, bias=True), nn.InstanceNorm2d(dim),
                   nn.ReLU(True)]

    conv_block += [nn.ReflectionPad2d(1)]
    conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=0, bias=True), nn.InstanceNorm2d(dim)]

    return nn.Sequential(*conv_block)


class ResNetBlock(nn.Module):  # <1> Create a model

    def __init__(self, dim):
        super(ResNetBlock, self).__init__()
        self.conv_block = build_conv_block(dim)

    def forward(self, x):
        out = x + self.conv_block(x)
        return out


class ResNetGenerator(nn.Module):

    def __init__(self, input_nc=3, output_nc=3, ngf=64, n_block=9):
        assert (n_block >= 0)
        super(ResNetGenerator, self).__init__()

        self.input_nc = input_nc
        self.output_nc = output_nc
        self.ngf = ngf

        model = [nn.ReflectionPad2d(3),
                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=True),
                 nn.InstanceNorm2d(ngf),
                 nn.ReLU(True)]

        n_downsampling = 2
        for i in range(n_downsampling):
            mult = 2 ** i
            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=True),
                      nn.InstanceNorm2d(ngf * mult * 2),
                      nn.ReLU(True)]
            mult = 2 ** n_downsampling

        for i in range(n_block):
            model += [ResNetBlock(ngf * mult)]

        for i in range(n_downsampling):
            mult = 2 ** (n_downsampling - i)
            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2,
                                         padding=1, output_padding=1, bias=True),
                      nn.InstanceNorm2d(ngf * mult / 2),
                      nn.ReLU(True)]

        model += [nn.ReflectionPad2d(3)]
        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
        model += [nn.Tanh()]

        self.model = nn.Sequential(*model)

    def forward(self, input):
        return self.model(input)


netG = ResNetGenerator()  # <2> Train a model
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
netG.to(device)
print(netG.to("cuda:0"))  # 打印出netG的网络结构
model_path = './data/Horse2Zebra/horse2zebra_0.4.0.pth'
model_data = torch.load(model_path, map_location=device)
netG.load_state_dict(model_data)  # 模型的加载

netG.eval()

preprocess = transforms.Compose([transforms.Resize(256),
                                 transforms.ToTensor()])

# 至此，模型准备完毕，开始加载图片进行测试
# 模型类别：基于ResNet的图像生成模型
# 实现功能：Horse2Zebra
# 原始图像:img
# 生成图像:out_img
img = Image.open("./data/Horse2Zebra/Beautiful-horse.jpg")
img_t = preprocess(img).unsqueeze(0).to(device)  # Move input tensor to the same device(关于输入模型张量与模型参数不匹配问题)
batch_out = netG(img_t)
out_t = (batch_out.data.squeeze() + 1.0) / 2.0
out_img = transforms.ToPILImage()(out_t)

plt.figure(figsize=(10, 5), facecolor='#1f5f8b')  # fig-size:指定figure的宽和高，单位为英寸
plt.subplot(1, 2, 1)
size = img.size
plt.title('Original Horse' + str(size), color='#a1de93')
plt.imshow(img)
plt.axis('off')  # 关掉坐标轴

plt.subplot(1, 2, 2)
size = img.size
plt.title('Generated Zebra' + str(size), color='#f38181')
plt.imshow(out_img)
plt.axis('off')

plt.show()
